import os
import sys
import argparse
import json
import base64
import time
from pathlib import Path
from typing import List, Tuple
import requests
from tqdm import tqdm
from PIL import Image

class PPTPhotoProcessor:
    def __init__(self, api_key: str, model_name: str = "anthropic/claude-3-sonnet"):
        """
        初始化PPT照片处理器
        
        Args:
            api_key: OpenRouter API密钥
            model_name: 要使用的多模态模型名称
        """
        self.api_key = api_key
        self.model_name = model_name
        self.base_url = "https://openrouter.ai/api/v1/chat/completions"
        self.headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
        
    def get_image_files(self, folder_path: str) -> List[str]:
        """
        获取文件夹中的所有图像文件
        
        Args:
            folder_path: 文件夹路径
            
        Returns:
            图像文件路径列表
        """
        image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
        image_files = []
        
        folder = Path(folder_path)
        if not folder.exists():
            raise FileNotFoundError(f"文件夹不存在: {folder_path}")
            
        for file_path in folder.iterdir():
            if file_path.suffix.lower() in image_extensions:
                image_files.append(str(file_path))
                
        # 按文件名排序
        image_files.sort()
        return image_files
    
    def encode_image_to_base64(self, image_path: str) -> str:
        """
        将图像编码为base64字符串
        
        Args:
            image_path: 图像文件路径
            
        Returns:
            base64编码的图像字符串
        """
        try:
            with Image.open(image_path) as img:
                # 如果图像过大，进行压缩
                max_size = (1024, 1024)
                if img.size[0] > max_size[0] or img.size[1] > max_size[1]:
                    img.thumbnail(max_size, Image.Resampling.LANCZOS)
                
                # 转换为RGB模式（如果需要）
                if img.mode != 'RGB':
                    img = img.convert('RGB')
                
                # 保存到内存中
                from io import BytesIO
                buffer = BytesIO()
                img.save(buffer, format='JPEG', quality=85)
                buffer.seek(0)
                
                # 编码为base64
                import base64
                return base64.b64encode(buffer.getvalue()).decode('utf-8')
                
        except Exception as e:
            raise Exception(f"处理图像 {image_path} 时出错: {str(e)}")
    
    def analyze_ppt_image(self, image_path: str, page_number: int) -> str:
        """
        分析单张PPT图像
        
        Args:
            image_path: 图像路径
            page_number: 页码
            
        Returns:
            分析结果文本
        """
        try:
            # 编码图像
            base64_image = self.encode_image_to_base64(image_path)
            
            # 构建请求消息
            messages = [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": """你是一个专业的PPT内容分析助手。请仔细分析这张PPT照片，并按照以下格式输出：

## 第X页

### 主要标题
[识别出的主要标题或主题]

### 关键内容
[列举主要的文字内容，保持原有的结构和逻辑]

### 要点总结
- [要点1]
- [要点2]
- [要点3]
...

请确保：
1. 准确识别图片中的所有可见文字
2. 保持原有的层次结构和逻辑关系
3. 如果有图表或示意图，请描述其主要信息
4. 忽略不重要的装饰性元素"""
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{base64_image}"
                            }
                        }
                    ]
                }
            ]
            
            # 发送API请求
            payload = {
                "model": self.model_name,
                "messages": messages,
                "max_tokens": 2000,
                "temperature": 0.1,
            }
            
            response = requests.post(
                self.base_url,
                headers=self.headers,
                json=payload,
                timeout=60
            )
            
            if response.status_code == 200:
                result = response.json()
                content = result['choices'][0]['message']['content']
                # 替换页码占位符
                content = content.replace("第X页", f"第{page_number}页")
                return content
            else:
                error_msg = f"API请求失败: {response.status_code} - {response.text}"
                print(f"警告: {error_msg}")
                return f"## 第{page_number}页\n\n### 处理失败\n{error_msg}\n\n"
                
        except Exception as e:
            error_msg = f"处理图像时出错: {str(e)}"
            print(f"警告: {error_msg}")
            return f"## 第{page_number}页\n\n### 处理失败\n{error_msg}\n\n"
    
    def generate_summary(self, content: str) -> str:
        """
        生成整体总结
        
        Args:
            content: 所有页面的内容
            
        Returns:
            总结文本
        """
        try:
            messages = [
                {
                    "role": "user",
                    "content": f"""请为以下PPT内容生成一个综合总结。总结应该包括：

1. 主要主题和议题
2. 关键观点和结论
3. 重要数据和事实
4. 整体结构概览

PPT内容：
{content}

请按照以下格式输出总结：

# PPT内容总结

## 主要主题

## 核心观点

## 重要信息

## 整体结构
"""
                }
            ]
            
            payload = {
                "model": self.model_name,
                "messages": messages,
                "max_tokens": 1500,
                "temperature": 0.2
            }
            
            response = requests.post(
                self.base_url,
                headers=self.headers,
                json=payload,
                timeout=60
            )
            
            if response.status_code == 200:
                result = response.json()
                return result['choices'][0]['message']['content']
            else:
                return "## 总结生成失败\n无法生成整体总结。\n\n"
                
        except Exception as e:
            return f"## 总结生成失败\n错误信息: {str(e)}\n\n"
    
    def process_folder(self, folder_path: str, output_path: str):
        """
        处理整个文件夹的PPT图像
        
        Args:
            folder_path: 输入文件夹路径
            output_path: 输出文件路径
        """
        # 获取所有图像文件
        image_files = self.get_image_files(folder_path)
        
        if not image_files:
            print("错误: 在指定文件夹中未找到图像文件")
            return
            
        print(f"找到 {len(image_files)} 个图像文件")
        
        # 创建输出目录
        output_file = Path(output_path)
        output_file.parent.mkdir(parents=True, exist_ok=True)
        
        # 处理每个图像文件
        all_content = []
        
        with tqdm(total=len(image_files), desc="处理PPT图像") as pbar:
            for i, image_path in enumerate(image_files, 1):
                pbar.set_description(f"处理第 {i} 页")
                
                # 分析图像
                content = self.analyze_ppt_image(image_path, i)
                all_content.append(content)
                
                # 增量式保存
                self.save_incremental_content(output_path, content)
                
                pbar.update(1)
                
                # 添加短暂延迟，避免API限流
                time.sleep(0.5)
        
        # 生成整体总结
        print("生成整体总结...")
        full_content = "\n\n".join(all_content)
        summary = self.generate_summary(full_content)
        
        # 保存最终文件（总结在前）
        self.save_final_content(output_path, summary, all_content)
        
        print(f"处理完成！结果已保存到: {output_path}")
    
    def save_incremental_content(self, output_path: str, content: str):
        """
        增量式保存内容
        
        Args:
            output_path: 输出文件路径
            content: 要追加的内容
        """
        with open(output_path, 'a', encoding='utf-8') as f:
            f.write(content + "\n\n")
    
    def save_final_content(self, output_path: str, summary: str, all_content: List[str]):
        """
        保存最终内容（总结在前）
        
        Args:
            output_path: 输出文件路径
            summary: 总结内容
            all_content: 所有页面内容
        """
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(summary + "\n\n")
            f.write("=" * 50 + "\n")
            f.write("# 详细内容\n\n")
            f.write("\n\n".join(all_content))


def main():
    parser = argparse.ArgumentParser(
        description="将手机拍摄的PPT照片转换为文字总结",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例用法:
  python ppt_processor.py -i ./ppt_photos -o ./output.txt -k your_api_key
  python ppt_processor.py --input-folder ./photos --output-file ./summary.txt --api-key your_key --model anthropic/claude-3-sonnet

支持的图像格式: JPG, JPEG, PNG, BMP, TIFF, WEBP
        """
    )
    
    parser.add_argument(
        '-i', '--input-folder',
        required=True,
        help='包含PPT照片的文件夹路径'
    )
    
    parser.add_argument(
        '-o', '--output-file',
        required=True,
        help='输出文本文件的路径'
    )
    
    parser.add_argument(
        '-k', '--api-key',
        help='OpenRouter API密钥（也可通过环境变量 OPENROUTER_API_KEY 设置）'
    )
    
    parser.add_argument(
        '-m', '--model',
        default='anthropic/claude-3-sonnet',
        help='要使用的模型名称 (默认: anthropic/claude-3-sonnet)'
    )
    
    args = parser.parse_args()
    
    # 获取API密钥
    api_key = args.api_key or os.getenv('OPENROUTER_API_KEY')
    if not api_key:
        print("错误: 请提供OpenRouter API密钥")
        print("方法1: 使用 -k 参数")
        print("方法2: 设置环境变量 OPENROUTER_API_KEY")
        sys.exit(1)
    
    # 验证输入路径
    if not os.path.exists(args.input_folder):
        print(f"错误: 输入文件夹不存在: {args.input_folder}")
        sys.exit(1)
    
    try:
        # 创建处理器并处理
        processor = PPTPhotoProcessor(api_key, args.model)
        processor.process_folder(args.input_folder, args.output_file)
        
    except KeyboardInterrupt:
        print("\n用户取消操作")
        sys.exit(0)
    except Exception as e:
        print(f"错误: {str(e)}")
        sys.exit(1)


if __name__ == "__main__":
    main()